﻿#include "SignatureUtils.h"
#include "Logger.h"

// 匿名空间
namespace {
	// 解析特征码字符串为字节序列和掩码序列
	static bool ParseSignatureString(const std::string& signatureStr, std::vector<BYTE>& outSignature, std::vector<BYTE>& outMask) {
		std::istringstream iss(signatureStr);
		std::string token;
		while (iss >> token) {
			if (token == "??" || token == "?") {
				outSignature.push_back(0x00);
				outMask.push_back(0x00);
			}
			else {
				try {
					BYTE val = static_cast<BYTE>(std::stoul(token, nullptr, 16));
					outSignature.push_back(val);
					outMask.push_back(0xFF);
				}
				catch (...) {
					return false;
				}
			}
		}
		return !outSignature.empty() && outSignature.size() == outMask.size();
	}

	// 掩码匹配函数：mask[i]==0xFF表示精确，mask[i]==0x00表示通配
	static bool MaskCompare(const BYTE* data, const BYTE* sig, const BYTE* mask, size_t len) {
		for (size_t i = 0; i < len; ++i) {
			if (mask[i] == 0xFF) {
				if (data[i] != sig[i]) return false;
			}
			// mask[i]==0x00表示通配，直接跳过
		}
		return true;
	}

	// 获取模块大小
	static size_t GetModuleSize(HMODULE hModule) {
		if (!hModule) return 0;

		PIMAGE_DOS_HEADER pDosHeader = reinterpret_cast<PIMAGE_DOS_HEADER>(hModule);
		if (pDosHeader->e_magic != IMAGE_DOS_SIGNATURE) return 0;

		PIMAGE_NT_HEADERS pNtHeaders = reinterpret_cast<PIMAGE_NT_HEADERS>(
			reinterpret_cast<BYTE*>(hModule) + pDosHeader->e_lfanew);
		if (pNtHeaders->Signature != IMAGE_NT_SIGNATURE) return 0;

		return pNtHeaders->OptionalHeader.SizeOfImage;
	}

	// 区段信息结构体
	struct SectionInfo {
		std::string name;
		uintptr_t start;
		uintptr_t end;
		DWORD characteristics;
	};

	// 获取模块的区段信息
	std::vector<SectionInfo> GetModuleSectionInfo(HMODULE hModule) {
		std::vector<SectionInfo> sections;

		PIMAGE_DOS_HEADER dosHeader = reinterpret_cast<PIMAGE_DOS_HEADER>(hModule);
		if (dosHeader->e_magic != IMAGE_DOS_SIGNATURE) {
			return sections;
		}

		PIMAGE_NT_HEADERS ntHeaders = reinterpret_cast<PIMAGE_NT_HEADERS>(
			reinterpret_cast<uintptr_t>(hModule) + dosHeader->e_lfanew);
		if (ntHeaders->Signature != IMAGE_NT_SIGNATURE) {
			return sections;
		}

		PIMAGE_SECTION_HEADER sectionHeader = IMAGE_FIRST_SECTION(ntHeaders);

		for (int i = 0; i < ntHeaders->FileHeader.NumberOfSections; i++) {
			SectionInfo info;

			// 获取区段名称
			char sectionName[9] = { 0 };
			memcpy(sectionName, sectionHeader[i].Name, 8);
			info.name = sectionName;

			// 计算区段起始和结束地址
			info.start = reinterpret_cast<uintptr_t>(hModule) + sectionHeader[i].VirtualAddress;
			info.end = info.start + sectionHeader[i].Misc.VirtualSize;
			info.characteristics = sectionHeader[i].Characteristics;

			sections.push_back(info);
		}

		return sections;
	}

	// 全模块扫描函数
	void* FindSignatureString(
		const char* moduleName,
		const std::string& signatureStr,
		bool searchDirection,
		size_t matchIndex
	) {
		// 解析特征码
		std::vector<BYTE> signature, mask;
		if (!ParseSignatureString(signatureStr, signature, mask)) {
			g_Logger.Log("[全模块扫描] 错误: 特征码字符串解析失败: %s", signatureStr.c_str());
			return nullptr;
		}

		HMODULE hModule = GetModuleHandleA(moduleName);
		if (!hModule) {
			g_Logger.Log("[全模块扫描] 错误: 无法获取模块句柄: %s", moduleName);
			return nullptr;
		}

		// 获取模块信息
		uintptr_t moduleBase = reinterpret_cast<uintptr_t>(hModule);
		size_t moduleSize = GetModuleSize(hModule);
		BYTE* moduleStart = reinterpret_cast<BYTE*>(hModule);
		BYTE* moduleEnd = moduleStart + moduleSize;

		// 获取所有区段信息
		auto sections = GetModuleSectionInfo(hModule);

		size_t currentMatch = 0;
		const size_t sigSize = signature.size();

		// 遍历整个模块内存
		for (BYTE* ptr = moduleStart; ptr < moduleEnd - sigSize; ptr++) {
			if (MaskCompare(ptr, signature.data(), mask.data(), sigSize)) {
				if (currentMatch == matchIndex) {
					// 确定所在区段
					std::string sectionName = "Unknown";
					uintptr_t addressValue = reinterpret_cast<uintptr_t>(ptr);

					for (const auto& section : sections) {
						if (addressValue >= section.start && addressValue < section.end) {
							sectionName = section.name;
							break;
						}
					}
					return ptr;
				}
				currentMatch++;
			}
		}

		return nullptr;
	}
};

namespace SignatureUtils {
	void* FindSignatureInSection(
		const char* moduleName,
		const std::vector<BYTE>& signature,
		const char* sectionName,
		bool searchDirection,
		size_t matchIndex)
	{
		g_Logger.Log("[特征码搜索] 开始在模块 %s 的 %s 段搜索特征码 (长度: %zu, 匹配序号: %zu)",
			moduleName, sectionName, signature.size(), matchIndex);

		if (signature.empty()) {
			g_Logger.Log("[特征码搜索] 错误: 特征码为空");
			return nullptr;
		}

		HMODULE hModule = GetModuleHandleA(moduleName);
		if (!hModule) {
			g_Logger.Log("[特征码搜索] 错误: 无法获取模块 %s 的句柄", moduleName);
			return nullptr;
		}

		PIMAGE_DOS_HEADER pDosHeader = reinterpret_cast<PIMAGE_DOS_HEADER>(hModule);
		if (pDosHeader->e_magic != IMAGE_DOS_SIGNATURE) {
			g_Logger.Log("[特征码搜索] 错误: 无效的DOS头");
			return nullptr;
		}

		PIMAGE_NT_HEADERS pNtHeaders = reinterpret_cast<PIMAGE_NT_HEADERS>(
			reinterpret_cast<BYTE*>(hModule) + pDosHeader->e_lfanew);
		if (pNtHeaders->Signature != IMAGE_NT_SIGNATURE) {
			g_Logger.Log("[特征码搜索] 错误: 无效的NT头");
			return nullptr;
		}

		PIMAGE_SECTION_HEADER pSection = IMAGE_FIRST_SECTION(pNtHeaders);
		for (WORD i = 0; i < pNtHeaders->FileHeader.NumberOfSections; ++i, ++pSection) {
			char currentSectionName[9] = { 0 };
			memcpy(currentSectionName, pSection->Name, 8);

			if (strncmp(currentSectionName, sectionName, 8) == 0) {
				BYTE* pStart = reinterpret_cast<BYTE*>(hModule) + pSection->VirtualAddress;
				size_t sectionSize = pSection->Misc.VirtualSize;
				if (sectionSize < signature.size()) {
					g_Logger.Log("[特征码搜索] 错误: 区段大小小于特征码长度");
					return nullptr;
				}
				size_t searchSize = sectionSize - signature.size() + 1;

				g_Logger.Log("[特征码搜索] 在 %s 段搜索 (范围: 0x%p - 0x%p, 大小: %zu)",
					sectionName, pStart, pStart + sectionSize, sectionSize);

				size_t currentMatch = 0;

				if (searchDirection) {
					for (size_t offset = 0; offset < searchSize; ++offset) {
						if (memcmp(pStart + offset, signature.data(), signature.size()) == 0) {
							if (currentMatch == matchIndex) {
								g_Logger.Log("[特征码搜索] 找到第 %zu 个匹配 (地址: 0x%p)",
									matchIndex + 1, pStart + offset);
								return pStart + offset;
							}
							currentMatch++;
						}
					}
				}
				else {
					for (size_t offset = searchSize; offset-- > 0;) {
						if (memcmp(pStart + offset, signature.data(), signature.size()) == 0) {
							if (currentMatch == matchIndex) {
								g_Logger.Log("[特征码搜索] 找到第 %zu 个匹配 (地址: 0x%p)",
									matchIndex + 1, pStart + offset);
								return pStart + offset;
							}
							currentMatch++;
						}
					}
				}
				g_Logger.Log("[特征码搜索] 共找到 %zu 个匹配项 (但请求的是第 %zu 个)",
					currentMatch, matchIndex + 1);
				break;
			}
		}

		g_Logger.Log("[特征码搜索] 未在 %s 段找到匹配特征码", sectionName);
		return nullptr;
	}

	void* FindSignatureWithMaskInSection(
		const char* moduleName,
		const std::vector<BYTE>& signature,
		const std::vector<BYTE>& mask,
		const char* sectionName,
		bool searchDirection,
		size_t matchIndex)
	{
		g_Logger.Log("[特征码搜索] 开始在模块 %s 的 %s 段搜索带掩码特征码 (长度: %zu, 匹配序号: %zu)",
			moduleName, sectionName, signature.size(), matchIndex);

		if (signature.empty() || mask.empty() || signature.size() != mask.size()) {
			g_Logger.Log("[特征码搜索] 错误: 特征码或掩码为空，或长度不一致");
			return nullptr;
		}

		HMODULE hModule = GetModuleHandleA(moduleName);
		if (!hModule) {
			g_Logger.Log("[特征码搜索] 错误: 无法获取模块 %s 的句柄", moduleName);
			return nullptr;
		}

		PIMAGE_DOS_HEADER pDosHeader = reinterpret_cast<PIMAGE_DOS_HEADER>(hModule);
		if (pDosHeader->e_magic != IMAGE_DOS_SIGNATURE) {
			g_Logger.Log("[特征码搜索] 错误: 无效的DOS头");
			return nullptr;
		}

		PIMAGE_NT_HEADERS pNtHeaders = reinterpret_cast<PIMAGE_NT_HEADERS>(
			reinterpret_cast<BYTE*>(hModule) + pDosHeader->e_lfanew);
		if (pNtHeaders->Signature != IMAGE_NT_SIGNATURE) {
			g_Logger.Log("[特征码搜索] 错误: 无效的NT头");
			return nullptr;
		}

		PIMAGE_SECTION_HEADER pSection = IMAGE_FIRST_SECTION(pNtHeaders);
		for (WORD i = 0; i < pNtHeaders->FileHeader.NumberOfSections; ++i, ++pSection) {
			char currentSectionName[9] = { 0 };
			memcpy(currentSectionName, pSection->Name, 8);

			if (strncmp(currentSectionName, sectionName, 8) == 0) {
				BYTE* pStart = reinterpret_cast<BYTE*>(hModule) + pSection->VirtualAddress;
				size_t sectionSize = pSection->Misc.VirtualSize;
				if (sectionSize < signature.size()) {
					g_Logger.Log("[特征码搜索] 错误: 区段大小小于特征码长度");
					return nullptr;
				}
				size_t searchSize = sectionSize - signature.size() + 1;

				g_Logger.Log("[特征码搜索] 在 %s 段搜索 (范围: 0x%p - 0x%p, 大小: %zu)",
					sectionName, pStart, pStart + sectionSize, sectionSize);

				size_t currentMatch = 0;

				if (searchDirection) {
					for (size_t offset = 0; offset < searchSize; ++offset) {
						if (MaskCompare(pStart + offset, signature.data(), mask.data(), signature.size())) {
							if (currentMatch == matchIndex) {
								g_Logger.Log("[特征码搜索] 找到第 %zu 个匹配 (地址: 0x%p)",
									matchIndex + 1, pStart + offset);
								return pStart + offset;
							}
							currentMatch++;
						}
					}
				}
				else {
					for (size_t offset = searchSize; offset-- > 0;) {
						if (MaskCompare(pStart + offset, signature.data(), mask.data(), signature.size())) {
							if (currentMatch == matchIndex) {
								g_Logger.Log("[特征码搜索] 找到第 %zu 个匹配 (地址: 0x%p)",
									matchIndex + 1, pStart + offset);
								return pStart + offset;
							}
							currentMatch++;
						}
					}
				}
				g_Logger.Log("[特征码搜索] 共找到 %zu 个匹配项 (但请求的是第 %zu 个)",
					currentMatch, matchIndex + 1);
				break;
			}
		}

		g_Logger.Log("[特征码搜索] 未在 %s 段找到匹配特征码", sectionName);
		return nullptr;
	}

	void* FindSignatureStringInSection(
		const char* moduleName,
		const std::string& signatureStr,
		const char* sectionName,
		bool searchDirection,
		size_t matchIndex)
	{
		std::vector<BYTE> signature, mask;
		if (!ParseSignatureString(signatureStr, signature, mask)) {
			g_Logger.Log("[特征码搜索] 错误: 特征码字符串解析失败: %s", signatureStr.c_str());
			return nullptr;
		}

		g_Logger.Log("[特征码搜索] 信息: 正在解析特征码模式: %s", signatureStr.c_str());
		return FindSignatureWithMaskInSection(moduleName, signature, mask, sectionName, searchDirection, matchIndex);
	}

	void* DevFindSignature(
		const char* moduleName,
		const std::string& signatureStr
	) {
		HMODULE hModule = GetModuleHandleA(moduleName);
		if (!hModule) {
			g_Logger.Log("[开发模式] 错误: 无法获取模块句柄: %s", moduleName);
			return nullptr;
		}

		// 获取模块信息
		uintptr_t moduleBase = reinterpret_cast<uintptr_t>(hModule);
		size_t moduleSize = GetModuleSize(hModule);

		// 获取所有区段信息
		auto sections = GetModuleSectionInfo(hModule);

		g_Logger.Log("\n==================================================");
		g_Logger.Log("[开发模式] 特征码搜索: %s", signatureStr.c_str());
		g_Logger.Log("模块: %s (基址: %p, 大小: 0x%X)", moduleName, hModule, moduleSize);
		g_Logger.Log("扫描模式: 全模块扫描");
		g_Logger.Log("==================================================");

		// 查找所有匹配项
		size_t matchCount = 0;
		void* firstMatch = nullptr;
		bool foundAtLeastOne = false;

		for (size_t i = 0; ; i++) {
			void* foundAddress = FindSignatureString(
				moduleName, signatureStr, true, i);

			if (!foundAddress) {
				// 只在真正找不到匹配项时打印
				if (i == 0) {
					g_Logger.Log("未找到匹配的特征码");
				}
				break;
			}

			foundAtLeastOne = true;  // 标记已找到至少一个匹配项

			// 记录第一个匹配项
			if (i == 0) firstMatch = foundAddress;

			// 确定特征码所在的区段
			std::string sectionName = "Unknown";
			uintptr_t addressValue = reinterpret_cast<uintptr_t>(foundAddress);

			for (const auto& section : sections) {
				if (addressValue >= section.start && addressValue < section.end) {
					sectionName = section.name;
					break;
				}
			}

			// 计算偏移量
			uintptr_t offset = addressValue - moduleBase;

			// 打印匹配信息（只在这里打印一次）
			g_Logger.Log("匹配项 #%d:", matchCount + 1);
			g_Logger.Log("  所在区段: %s", sectionName.c_str());
			g_Logger.Log("  绝对地址: %p", foundAddress);
			g_Logger.Log("  相对偏移: 0x%X", offset);
			g_Logger.Log("  发布接口调用: GetHardcodedAddress(\"%s\", 0x%X, 0x%X)",
				moduleName, offset, moduleSize);

			matchCount++;

			// 安全限制：最多100个匹配项
			if (matchCount >= 100) {
				g_Logger.Log("警告: 已达到最大匹配项限制(100)");
				break;
			}
		}

		if (matchCount == 0) {
			g_Logger.Log("==================================================\n");
			return nullptr;
		}

		g_Logger.Log("==================================================");
		g_Logger.Log("总计匹配: %d 个特征码", matchCount);
		if (firstMatch) {
			g_Logger.Log("第一个匹配地址: %p (可选返回)", firstMatch);
		}
		g_Logger.Log("==================================================\n");

		return firstMatch;
	}

	void* GetHardcodedAddress(
		const char* moduleName,
		uintptr_t hardcodedOffset,
		size_t expectedModuleSize
	) {
		// 1. 获取模块句柄
		HMODULE hModule = GetModuleHandleA(moduleName);
		if (!hModule) {
			// 发布模式下使用Windows API弹窗
			char msg[256];
			sprintf_s(msg, "模块加载失败: %s", moduleName);
			MessageBoxA(NULL, msg, "特征码定位错误", MB_ICONERROR);
			return nullptr;
		}

		// 2. 校验模块大小
		size_t actualModuleSize = GetModuleSize(hModule);
		if (actualModuleSize != expectedModuleSize) {
			char msg[256];
			sprintf_s(msg,
				"模块大小不匹配!\n模块: %s\n预期: 0x%zX\n实际: 0x%zX",
				moduleName, expectedModuleSize, actualModuleSize);
			MessageBoxA(NULL, msg, "特征码定位错误", MB_ICONERROR);
			return nullptr;
		}

		// 3. 计算最终地址
		return reinterpret_cast<void*>(
			reinterpret_cast<uintptr_t>(hModule) + hardcodedOffset);
	}
};